import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.utils.data.sampler import BatchSampler, SubsetRandomSampler
from torch.distributions import Categorical
from masking_categorical import CategoricalMasked
import config
from object import delay
import random
import math
import os

import pandas as pd
import dcn_4 as dcn
from sklearn.preprocessing import LabelEncoder, MinMaxScaler
from deepctr_torch.models import xDeepFM
from deepctr_torch.inputs import SparseFeat, DenseFeat, get_feature_names

gpus = [0]
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Trick 8: orthogonal initialization
def orthogonal_init(layer, gain=1.0):
    nn.init.orthogonal_(layer.weight, gain=gain)
    nn.init.constant_(layer.bias, 0)


class Actor(nn.Module):
    def __init__(self, args):
        super(Actor, self).__init__()
        
        self.fc1 = nn.Linear((config.EDGE_NODE_NUM+1)*32, 512)
        self.fc2 = nn.Linear(512,128)
        self.fc3 = nn.Linear(128, 128)
        self.fc4 = nn.Linear(128, args.action_dim)
        
        self.activate_func = [nn.ReLU(), nn.Tanh()][args.use_tanh]
        self.deepdcn = None
        self.embedding_size = 64

        # parameters of DCN
        self.dense_features = ['C' + str(i) for i in range(1, 13)]
        self.sparse_features = ['I' + str(i) for i in range(1, 720)]

        if args.use_orthogonal_init:
            print("------use_orthogonal_init------")
            orthogonal_init(self.fc1)
            orthogonal_init(self.fc2)
            orthogonal_init(self.fc3)
            orthogonal_init(self.fc4, gain=0.01)
            # orthogonal_init(self.fc5, gain=0.01)

    def forward(self, s):
        if self.deepdcn is None:

            s_1 = s.cpu().detach().numpy()
            # s_.resize((config.EDGE_NODE_NUM+1,731),refcheck=False)
            s_ = np.resize(s_1, (config.EDGE_NODE_NUM+1,731))

            self.s_df = pd.DataFrame(s_, columns = [self.dense_features + self.sparse_features])
            
            # print("s_ is: ", self.s_df)

            feat_sizes = {}
            feat_sizes_dense = {feat: 1 for feat in self.dense_features}
            feat_sizes_sparse = {feat: len(self.s_df[feat]) for feat in self.sparse_features}
            feat_sizes.update(feat_sizes_dense)
            feat_sizes.update(feat_sizes_sparse)

            # exit(1)
            for feat in self.sparse_features:
                lbe = LabelEncoder()
                self.s_df[feat] = lbe.fit_transform(self.s_df[feat])

            
            mms = MinMaxScaler(feature_range=(0, 1))
            self.s_df[self.dense_features] = mms.fit_transform(self.s_df[self.dense_features])

            # self.fixlen_feature_columns = [SparseFeat(feat, self.s_df[feat].nunique()[0]) for feat in self.sparse_features] + [DenseFeat(feat, 1, ) for feat in self.dense_features]
            self.fixlen_feature_columns = [(feat, 'sparse') for feat in self.sparse_features] + [(feat, 'dense') for feat in
                                                                              self.dense_features]
            # print(self.fixlen_feature_columns)
            self.dnn_feature_columns = self.fixlen_feature_columns
            self.linear_feature_columns = self.fixlen_feature_columns
            
            # self.deepdcn = xDeepFM(self.linear_feature_columns,self.dnn_feature_columns,task='regression',device=device)
            self.deepdcn = dcn.DCN(feat_sizes, self.embedding_size, self.linear_feature_columns, self.dnn_feature_columns).to(device)
        
        s_dcn = self.activate_func(self.deepdcn(s)).squeeze(1)
        
        s1 = self.activate_func(self.fc1(s_dcn))
        s2 = self.activate_func(self.fc2(s1))
        s3 = self.activate_func(self.fc3(s2))
        
        a_prob = torch.softmax(self.fc4(s3), dim=1)
        # a_prob = torch.softmax(s_dcn, dim=1).squeeze(1)
        
        return a_prob


class Critic(nn.Module):
    def __init__(self, args):
        super(Critic, self).__init__()
        self.fc1 = nn.Linear(args.statefinal_dim, args.hidden_width)
        self.fc2 = nn.Linear(args.hidden_width, args.hidden_width)
        self.fc3 = nn.Linear(args.hidden_width, 1)
        self.activate_func = [nn.ReLU(), nn.Tanh()][args.use_tanh]  

        if args.use_orthogonal_init:
            print("------use_orthogonal_init------")
            orthogonal_init(self.fc1)
            orthogonal_init(self.fc2)
            orthogonal_init(self.fc3)

    def forward(self, s):
        s = torch.flatten(s, start_dim=1)
        
        s = self.activate_func(self.fc1(s))
        s = self.activate_func(self.fc2(s))
        v_s = self.fc3(s)
        return v_s

class RunningMeanStd:
    # Dynamically calculate mean and std
    def __init__(self, shape):  # shape:the dimension of input data
        self.n = 0
        self.mean = np.zeros(shape)
        self.S = np.zeros(shape)
        self.std = np.sqrt(self.S)

    def update(self, x):
        x = np.array(x)
        self.n += 1
        if self.n == 1:
            self.mean = x
            self.std = x
        else:
            old_mean = self.mean.copy()
            self.mean = old_mean + (x - old_mean) / self.n
            self.S = self.S + (x - old_mean) * (x - self.mean)
            self.std = np.sqrt(self.S / self.n)


class RewardScaling:
    def __init__(self, shape, gamma):
        self.shape = shape  # reward shape=1
        self.gamma = gamma  # discount factor
        self.running_ms = RunningMeanStd(shape=self.shape)
        self.R = np.zeros(self.shape)

    def __call__(self, x):
        self.R = self.gamma * self.R + x
        self.running_ms.update(self.R)
        x = x / (self.running_ms.std + 1e-8)  # Only divided std
        return x

    def reset(self):  # When an episode is done,we should reset 'self.R'
        self.R = np.zeros(self.shape)


class PPO_discrete:
    def __init__(self, args):
        self.batch_size = args.batch_size
        self.batch_size_pre = args.batch_size_pre
        self.mini_batch_size = args.mini_batch_size
        self.max_train_steps = args.max_train_steps
        self.lr = args.lr
        self.lr_a = args.lr_a  # Learning rate of actor
        self.lr_c = args.lr_c  # Learning rate of critic
        self.gamma = args.gamma  # Discount factor
        self.lamda = args.lamda  # GAE parameter
        self.epsilon = args.epsilon  # PPO clip parameter
        self.K_epochs = args.K_epochs  # PPO parameter
        self.entropy_coef = args.entropy_coef  # Entropy coefficient
        self.set_adam_eps = args.set_adam_eps
        self.use_grad_clip = args.use_grad_clip
        self.use_lr_decay = args.use_lr_decay
        self.use_adv_norm = args.use_adv_norm

        self.actor = Actor(args).to(device)
        self.critic = Critic(args).to(device)
        self.rewardscaling = RewardScaling(1, self.gamma)

        if self.set_adam_eps:  
            self.optimizer_actor = torch.optim.Adam(self.actor.parameters(), lr=self.lr_a, eps=1e-5)
            self.optimizer_critic = torch.optim.Adam(self.critic.parameters(), lr=self.lr_c, eps=1e-5)
        else:
            self.optimizer_actor = torch.optim.Adam(self.actor.parameters(), lr=self.lr_a)
            self.optimizer_critic = torch.optim.Adam(self.critic.parameters(), lr=self.lr_c)

    def choose_action(self, s,uid,obs):
        available_actions = obs['next_can']
        ttt = []
        for id,item in enumerate(available_actions):
            if item == 1:
                ttt.append(id)
        # print('(choose_action)available_actions: ',ttt)

        s = torch.unsqueeze(torch.tensor(s, dtype=torch.float), 0).to(device)
        # exit(39)
        with torch.no_grad():
            available_actions = torch.tensor(available_actions,dtype=torch.bool)
            probs=self.actor(s)
            # print("probs is : ",probs)
            dist = CategoricalMasked(logits=probs,mask=available_actions)
            # a = dist.sample()
            a = dist.argmax()
        if len(ttt) != 0:
            action_retry_count = 0
            usr_locx = [j for i,j,k in obs['each_use_loc'] if i == uid][0]
            usr_locy = [k for i,j,k in obs['each_use_loc'] if i == uid][0]
        
            while self.satisfied_constrain(a.item(),uid,obs) == False: # not satify constrains
                a = dist.sample() # arbitrary prob distribution
                action_retry_count += 1
                if action_retry_count >= config.EDGE_NODE_NUM or usr_locx==np.inf or usr_locy==np.inf:
                    # raise EnvironmentError
                    a = config.EDGE_NODE_NUM # to cloud
                    a = torch.tensor([a])
                    break
        else:
            a = config.EDGE_NODE_NUM
            a = torch.tensor([a])

        with torch.no_grad():
            a = a.to(device)
            a_logprob = dist.log_prob(a)

        a = a.item()
        a_logprob = a_logprob.item()
        # return a.numpy()[0], a_logprob.numpy()[0]
        return a, a_logprob
    
    def choose_action_2(self, s,uid,obs,target):
        available_actions = obs['next_can']
        ttt = []
        for id,item in enumerate(available_actions):
            if item == 1:
                ttt.append(id)
        # print('(choose_action)available_actions: ',ttt)

        s = torch.unsqueeze(torch.tensor(s, dtype=torch.float), 0).to(device)
        # exit(39)
        with torch.no_grad():
            available_actions = torch.tensor(available_actions,dtype=torch.bool)
            probs=self.actor(s)
            dist = CategoricalMasked(logits=probs,mask=available_actions)
            # a = dist.sample()
            # a = dist.argmax()
            
        a = target
        a = torch.tensor([a])

        with torch.no_grad():
            a = a.to(device)
            a_logprob = dist.log_prob(a)

        a = a.item()
        a_logprob = a_logprob.item()
        # return a.numpy()[0], a_logprob.numpy()[0]
        return a, a_logprob
    
    def satisfied_constrain(self, edge_id, uidre, obs=None):
        # original from baseline.py
        # print("each edge_id is: ", edge_id)
        one_ulyr = [j for i,j in obs['each_usrhas_lyer'] if i == uidre][0]
        if edge_id == config.EDGE_NODE_NUM:
            return True
        edge_con = [d for a,b,c,d,f in obs['each_edge_cpumemcondisk'] if a == edge_id][0]
        if config.node_max_container_number - edge_con -1 < 0:
            print('edge {} container num is {} > 5.'.format(edge_id,edge_con))
            return  False
        # edge_usrnum_limit
        if config.node_usrnum_limit - obs['each_edge_usrnum'][edge_id] -1 < 0:
            print('edge {} has user num {} > 50.'.format(edge_id,obs['each_edge_usrnum'][edge_id]))
            return  False
        # edge_storage_free - task_size - download_size
        tmp_downsize = 0
        for item in one_ulyr:
            if item not in obs['each_edgehas_lyr'][edge_id]:
                tmp_downsize += obs['all_layer_downsiz'][item]

        edge_disk = [f for a,b,c,d,f in obs['each_edge_cpumemcondisk'] if a == edge_id][0]

        usr_has_tsksiz = [k for i,j,k in obs['each_usrhas_task'] if i == uidre][0]

        # print("edge {a} has disk {b} - layer_download_size:{d}".format(a=edge_id, b=edge_disk, d=tmp_downsize))
        if edge_disk - usr_has_tsksiz - tmp_downsize <0:
            print('edge {a} has disk {b} - user_tsk_size:{c} - layer_download_size:{d} < 0.'.format(a=edge_id,b=edge_disk,c=usr_has_tsksiz,d=tmp_downsize))
            return  False

    def update(self, replay_buffer, total_steps):
        s, a, a_logprob, r, s_, dw, done = replay_buffer.numpy_to_tensor()
        """
            Calculate the advantage using GAE
            'dw=True' means dead or win, there is no next state s'
            'done=True' represents the terminal of an episode(dead or win or reaching the max_episode_steps). When calculating the adv, if done=True, gae=0
        """
        a = a.to(device)
        s_ = s_.to(device)
        r = r.to(device)
        s = s.to(device)
        dw = dw.to(device)
        a_logprob = a_logprob.to(device)
        adv = []
        gae = 0
        with torch.no_grad():
            a = a.to(device)
            s_ = s_.to(device)
            r = r.to(device)
            s = s.to(device)
            dw = dw.to(device)
            a_logprob = a_logprob.to(device)

            vs = self.critic(s)
            vs_ = self.critic(s_)
            deltas = r + self.gamma * (1.0 - dw) * vs_ - vs
            for delta, d in zip(reversed(deltas.flatten().cpu().numpy()), reversed(done.flatten().cpu().numpy())):
                gae = delta + self.gamma * gae * (1.0 - d)
                adv.insert(0,gae)
            adv = torch.tensor(adv, dtype=torch.float).view(-1,1).to(device)
            v_target = adv + vs
            if self.use_adv_norm:
                adv = ((adv - adv.mean()) / (adv.std() + 1e-5))

        for _ in range(self.K_epochs):
            for index in BatchSampler(SubsetRandomSampler(range(self.batch_size)), self.mini_batch_size, False):
                dist_now = Categorical(probs=self.actor(s[index]))
                dist_entropy = dist_now.entropy().view(-1,1)
                a_logprob_now = dist_now.log_prob(a[index].squeeze()).view(-1,1).to(device)
                radios = torch.exp(a_logprob_now - a_logprob[index])

                surr1 = radios * adv[index]
                surr2 = torch.clamp(radios, 1 - self.epsilon, 1 + self.epsilon) * adv[index]
                actor_loss = -torch.min(surr1,surr2) - self.entropy_coef * dist_entropy

                # self.optimizer_actor.zero_grad()
                actor_loss = actor_loss.mean()
                self.optimizer_actor.zero_grad()
                actor_loss.backward()
                if self.use_grad_clip:
                    torch.nn.utils.clip_grad_norm_(self.actor.parameters(), 0.5)
                self.optimizer_actor.step()

                v_s = self.critic(s[index])
                critic_loss = F.mse_loss(v_target[index], v_s)

                self.optimizer_critic.zero_grad()
                critic_loss.backward()
                if self.use_grad_clip:
                    torch.nn.utils.clip_grad_norm_(self.critic.parameters(), 0.5)
                self.optimizer_critic.step()

        if self.use_lr_decay:
            self.lr_decay(total_steps)

        return actor_loss, critic_loss
    
    def update1(self, replay_buffer, total_steps):
        s, a, a_logprob, r, s_, dw, done = replay_buffer.numpy_to_tensor()
        """
            Calculate the advantage using GAE
            'dw=True' means dead or win, there is no next state s'
            'done=True' represents the terminal of an episode(dead or win or reaching the max_episode_steps). When calculating the adv, if done=True, gae=0
        """
        a = a.to(device)
        s_ = s_.to(device)
        r = r.to(device)
        s = s.to(device)
        dw = dw.to(device)
        a_logprob = a_logprob.to(device)
        adv = []
        gae = 0
        with torch.no_grad():
            a = a.to(device)
            s_ = s_.to(device)
            r = r.to(device)
            s = s.to(device)
            dw = dw.to(device)
            a_logprob = a_logprob.to(device)

            vs = self.critic(s)
            vs_ = self.critic(s_)
            deltas = r + self.gamma * (1.0 - dw) * vs_ - vs
            for delta, d in zip(reversed(deltas.flatten().cpu().numpy()), reversed(done.flatten().cpu().numpy())):
                gae = delta + self.gamma * gae * (1.0 - d)
                adv.insert(0,gae)
            adv = torch.tensor(adv, dtype=torch.float).view(-1,1).to(device)
            v_target = adv + vs
            if self.use_adv_norm:
                adv = ((adv - adv.mean()) / (adv.std() + 1e-5))

        for _ in range(self.K_epochs):
            for index in BatchSampler(SubsetRandomSampler(range(self.batch_size_pre)), self.mini_batch_size, False):
                # print("s[index] is : ",s[index])
                # print(len(s[index]))
                # print(len(s))
                # print(len(s[index][0]))
                dist_now = Categorical(probs=self.actor(s[index]))
                dist_entropy = dist_now.entropy().view(-1,1)
                a_logprob_now = dist_now.log_prob(a[index].squeeze()).view(-1,1).to(device)
                radios = torch.exp(a_logprob_now - a_logprob[index])

                surr1 = radios * adv[index]
                surr2 = torch.clamp(radios, 1 - self.epsilon, 1 + self.epsilon) * adv[index]
                actor_loss = -torch.min(surr1,surr2) - self.entropy_coef * dist_entropy

                # self.optimizer_actor.zero_grad()
                actor_loss = actor_loss.mean()
                self.optimizer_actor.zero_grad()
                actor_loss.backward()
                if self.use_grad_clip:
                    torch.nn.utils.clip_grad_norm_(self.actor.parameters(), 0.5)
                self.optimizer_actor.step()

                v_s = self.critic(s[index])
                critic_loss = F.mse_loss(v_target[index], v_s)

                self.optimizer_critic.zero_grad()
                critic_loss.backward()
                if self.use_grad_clip:
                    torch.nn.utils.clip_grad_norm_(self.critic.parameters(), 0.5)
                self.optimizer_critic.step()

        if self.use_lr_decay:
            self.lr_decay(total_steps)

        return actor_loss, critic_loss

    def lr_decay(self, total_steps):
        
        if total_steps == 150:
            
            lr_a_now = 3e-4
            lr_c_now = 3e-4
            lr_now = 3e-4
            
        else:
            lr_a_now = self.lr_a * (1 - total_steps / self.max_train_steps)
            lr_c_now = self.lr_c * (1 - total_steps / self.max_train_steps)
            lr_now = self.lr * (1 - total_steps / self.max_train_steps)
        
        for p in self.optimizer_actor.param_groups:
            p['lr'] = lr_a_now
        for p in self.optimizer_critic.param_groups:
            p['lr'] = lr_c_now
        for p in self.optimizer_actor.param_groups:
            p['lr'] = lr_now

    def lr_decay1(self,total_steps):
        lr_a_now = self.lr_a * math.pow(0.95,total_steps/45)
        lr_c_now = self.lr_c * math.pow(0.95,total_steps/45)
        lr_now = self.lr * math.pow(0.98,total_steps/45)
        for p in self.optimizer_actor.param_groups:
            p['lr'] = lr_a_now
        for p in self.optimizer_critic.param_groups:
            p['lr'] = lr_c_now
        for p in self.optimizer_actor.param_groups:
            p['lr'] = lr_now

    def save_model(self, total_epochs,path):
        if not os.path.exists(path):
            open(path, 'a').close()
        checkpoint = {
            "net_actor": self.actor.state_dict(),
            "net_ctritic": self.critic.state_dict(),
            "optimizer": self.optimizer_actor.state_dict(),
            "optimizer": self.optimizer_critic.state_dict(),
            "epoch": total_epochs,
        }
        torch.save(checkpoint, path)

    def load_model(self, path):
        checkpoint = torch.load(path)
        
        # self.ac.load_state_dict(checkpoint["net"])
        self.actor.load_state_dict(checkpoint["net"])
        self.critic.load_state_dict(checkpoint["net"])
        self.optimizer_actor.load_state_dict(checkpoint["optimizer"])
        epoch = checkpoint["epoch"]
        return epoch
